import unittest

from psychsim.pwl import *
from psychsim.action import Action,ActionSet
from psychsim.world import World,stateKey
from psychsim.agent import Agent
from psychsim.reward import *

''' 
This example is used to test changing the beliefs of an agent
depending on the dynamics of an action
'''

class LightSwitchTestCase(unittest.TestCase):

    def setUp(self):
        #setup world, agent, and basics
        self.maxRounds=5
        self.world = World()
        self.blindperson = Agent('BlindPerson')
        self.world.addAgent(self.blindperson)
        
        self.world.setOrder(['BlindPerson'])
                    
    def addStates(self):
        #define states            
        self.world.defineState(None,'round',int)
        self.world.setState(None,'round',0)
        
        self.world.defineState(None,'light',bool)
        self.world.setState(None,'light',True)
                  
    def addTermination(self):
        #terminate if maxRounds reached
        self.world.addTermination(makeTree({'if': thresholdRow(stateKey(None,'round'),self.maxRounds),
                                   True: True, False: False}))
           
    def addActions(self):
        #add actions and appropriate dynamics
        action =  self.blindperson.addAction({'verb': 'do nothing'})
        tree = makeTree(incrementMatrix(stateKey(None,'round'),1))
        self.world.setDynamics('round',action,tree) 
        
        action =  self.blindperson.addAction({'verb': 'action switch'})
        tree = makeTree(incrementMatrix(stateKey(None,'round'),1))
        self.world.setDynamics('round',action,tree)
        tree = makeTree({'if': trueRow(stateKey(None, 'light')),
                         True: setFalseMatrix(stateKey(None, 'light')),
                         False: setTrueMatrix(stateKey(None, 'light'))})
        self.world.setDynamics('light',action,tree)

    def setModelRewardsBeliefs(self):
        #set belief, rewards, and model to update                 
        key = stateKey(None,'light')
        self.blindperson.setBelief(key,Distribution({0: 1., 1: 0.}))
        
        self.blindperson.setReward(maximizeFeature(stateKey(None,'light')), 1., True)
        
        self.world.setModel('BlindPerson',True)

    def testSetup(self):
        self.addStates()
        self.addTermination()
        self.addActions()
        self.setModelRewardsBeliefs()           
            
    def testRun(self):
        self.testSetup()
        self.world.printState()
        self.world.printModels()
        for t in range(self.maxRounds + 1):
            self.world.explain(self.world.step(),level=1)
            self.world.printState()
            self.world.printModels()
            if self.world.terminated():
                break 
            
if __name__ == "__main__":
    suite = unittest.TestSuite()
    for method in dir(LightSwitchTestCase):
       if method.startswith("test"):
          suite.addTest(LightSwitchTestCase(method))
    unittest.TextTestRunner().run(suite)